# 功能说明：多轮对话，支持 OpenAI 和 DeepSeek 两种 API 调用，支持流式输出。
import requests
import openai
import os
from typing import List, Dict, Optional, Literal, Union

class UnifiedChatBot:
    def __init__(
        self,
        api_key: str,
        api_provider: Literal["openai", "deepseek"] = "openai",
        model: str = "gpt-3.5-turbo",  # 默认 OpenAI 模型
        system_prompt: Optional[str] = None,
        max_history: int = 10,
    ):
        """
        初始化统一聊天机器人（兼容 OpenAI 和 DeepSeek）
        
        参数:
            api_key: API 密钥
            api_provider: 服务商 ("openai" 或 "deepseek")
            model: 模型名称（如 "gpt-4"/"deepseek-chat"）
            system_prompt: 系统角色设定
            max_history: 最大对话历史长度
        """
        self.api_key = api_key
        self.api_provider = api_provider
        self.model = model
        self.max_history = max_history
        self.conversation_history: List[Dict] = []
        
        # 设置 API 终结点
        self.api_url = (
            "https://api.deepseek.com/v1/chat/completions" 
            if api_provider == "deepseek" 
            else None  # OpenAI 使用官方库，无需 URL
        )
        
        # 初始化系统提示
        if system_prompt:
            self.add_message("system", system_prompt)

    def _trim_history(self):
        """修剪对话历史，避免超出长度限制"""
        if len(self.conversation_history) > self.max_history:
            # 保留 system 消息（如果存在）
            system_msg = (
                [self.conversation_history[0]] 
                if self.conversation_history[0]["role"] == "system" 
                else []
            )
            self.conversation_history = system_msg + self.conversation_history[-(self.max_history - len(system_msg)):]

    def add_message(self, role: str, content: str):
        """添加消息到历史记录"""
        self.conversation_history.append({"role": role, "content": content})
        self._trim_history()

    def _call_openai(self, stream: bool = False) -> Union[str, None]:
        """调用 OpenAI API"""
        openai.api_key = self.api_key
        response = openai.ChatCompletion.create(
            model=self.model,
            messages=self.conversation_history,
            stream=stream,
        )
        
        if stream:
            full_response = ""
            print("AI: ", end="", flush=True)
            for chunk in response:
                content = chunk.choices[0].delta.get("content", "")
                print(content, end="", flush=True)
                full_response += content
            print()
            return full_response
        else:
            return response.choices[0].message.content

    def _call_deepseek(self, stream: bool = False) -> Union[str, None]:
        """调用 DeepSeek API"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": self.model,
            "messages": self.conversation_history,
            "stream": stream,
        }
        
        response = requests.post(
            self.api_url,
            headers=headers,
            json=payload,
            stream=stream,
        )
        
        if response.status_code != 200:
            error_msg = response.json().get("error", {}).get("message", "Unknown error")
            raise Exception(f"{self.api_provider} API Error: {error_msg}")
        
        if stream:
            full_response = ""
            print("AI: ", end="", flush=True)
            for chunk in response.iter_lines():
                if chunk:
                    decoded_chunk = chunk.decode("utf-8")
                    if decoded_chunk.startswith("data: "):
                        data = decoded_chunk[6:]
                        if data != "[DONE]":
                            try:
                                chunk_data = json.loads(data)
                                delta = chunk_data["choices"][0]["delta"]
                                if "content" in delta:
                                    content = delta["content"]
                                    print(content, end="", flush=True)
                                    full_response += content
                            except json.JSONDecodeError:
                                continue
            print()
            return full_response
        else:
            return response.json()["choices"][0]["message"]["content"]

    def get_response(self, user_input: str, stream: bool = False) -> str:
        """
        获取 AI 回复（自动根据 api_provider 选择接口）
        
        参数:
            user_input: 用户输入文本
            stream: 是否流式输出
            
        返回:
            AI 生成的回复内容
        """
        self.add_message("user", user_input)
        
        try:
            if self.api_provider == "openai":
                ai_response = self._call_openai(stream)
            else:
                ai_response = self._call_deepseek(stream)
            
            if ai_response:
                self.add_message("assistant", ai_response)
            return ai_response or ""
        
        except Exception as e:
            self.conversation_history.pop()  # 移除失败的用户输入
            raise e


# ===== 使用示例 =====
if __name__ == "__main__":
    # 配置选择
    API_PROVIDER = "deepseek"  # 可选 "openai" 或 "deepseek"
    API_KEY = api_key=os.getenv("OPENAI_API_KEY")
    MODEL = "deepseek-chat" if API_PROVIDER == "deepseek" else "gpt-3.5-turbo"

    # 初始化机器人
    bot = UnifiedChatBot(
        api_key=API_KEY,
        api_provider=API_PROVIDER,
        model=MODEL,
        system_prompt="你是一个专业且友好的助手。",
    )

    print(f"{API_PROVIDER.upper()} 聊天机器人已启动，输入 'exit' 退出\n")

    while True:
        user_input = input("你: ")
        if user_input.lower() in ["exit", "quit"]:
            break
        
        try:
            # 获取回复（stream=False 可改为 True 启用流式输出）
            ai_response = bot.get_response(user_input, stream=False)
            print(f"AI: {ai_response}\n")
        except Exception as e:
            print(f"错误: {str(e)}")